# _*_ coding: utf-8 _*_
import numpy as np
import matplotlib.pyplot as plt
from copy import deepcopy
from random import shuffle, uniform
from tqdm import trange

from VCMTSP import Instance

np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)
plt.rc('font', family='Times New Roman', weight='normal', size=12)


class ITPLS:
    def __init__(self, data):
        # -- read instance from instance data --
        self.V = data.cities  # city list
        self.n = len(self.V)  # the number of cities
        self.m = data.m  # number of salesmen
        self.A = data.A  # accessibility matrix
        self.Depot = data.depot  # depot
        self.V.remove(self.Depot)  # exclude depot
        self.Distance = data.distance_matrix  # distance matrix
        # -- parameters ---
        self.Ps = 0.5  # decision of intra-routing optimization
        self.alpha = 0.5
        self.O_max = 100  # search depth of SbTS
        self.Tl = 0.3  # tenure
        self.Opt_depth = 100  # 2-opt maximal depth
        self.L = int(1e8)  # the size of hash table
        self.T = 100
        # -- solution initialization --
        self.solution = self.greedy_init()
        self.Best = deepcopy(self.solution)
        self.evo_list = []

    def optimize(self, iteration=1000, trial=1):
        pbar = trange(iteration)
        pbar.set_description(f'{self.__class__.__name__} trial {trial}')
        ni = 0  # not improved iterations
        for _ in pbar:
            s_best = self.local_optima_exploration(s=self.solution)
            if s_best.total < self.Best.total:
                self.Best = deepcopy(s_best)
                ni = 0
            else:
                ni = ni + 1
            self.solution = self.local_optima_escaping(s=self.solution, ni=ni)
            self.evo_list.append(self.Best.total)
            pbar.set_postfix_str(f'best: {self.Best.total:.2f}')
        return self.Best

    class Solution:
        def __init__(self):
            self.tours = None
            self.costs = None
            self.total = 1e99

    def greedy_init(self):
        solution = self.Solution()
        solution.tours = [[self.Depot, self.Depot] for _ in range(self.m)]
        V = deepcopy(self.V)
        shuffle(V)
        for city in V:
            self.optimal_insert(solution, city)
        return self.calc_length(solution)

    def optimal_insert(self, solution, city, salesman=None, execute=True):
        """
        @param solution: [Solution] the solution that performs the optimal insertion
        @param city: [int] the city to be inserted
        @param salesman: [int] assign a specific salesman to insert the city
        @param execute: [bool] whether performs insertion in place
        @return: solution or optimal insertion cost
        """
        assert solution.tours is not None
        salesmen = self.salesmen_of(city) if salesman is None else [salesman]
        opt_r, opt_i, opt_cost = -1, -1, 1e99
        for r in salesmen:
            for i in range(1, len(solution.tours[r])):
                p, q = solution.tours[r][i - 1], solution.tours[r][i]
                cost = self.Distance[p][city] + self.Distance[city][q] - self.Distance[p][q]
                if opt_cost > cost:
                    opt_r, opt_i, opt_cost = r, i, cost
        if execute:
            solution.tours[opt_r].insert(opt_i, city)
            return solution
        else:
            return opt_cost

    def calc_length(self, solution):
        solution.costs = [0 for _ in solution.tours]
        for r, tour in enumerate(solution.tours):
            for i in range(1, len(tour)):
                solution.costs[r] += self.Distance[tour[i - 1]][tour[i]]
        solution.total = sum(solution.costs)
        return solution

    def salesmen_of(self, city):
        return np.where(self.A[city, :] == 1)[0].tolist()

    def cities_of(self, salesman):
        return np.where(self.A[:, salesman] == 1)[0].tolist()

    def local_optima_exploration(self, s):
        s = self.intra_route_optimization(s)
        s_best = deepcopy(s)
        R = set()
        H1, H2 = np.zeros(self.L, dtype=int), np.zeros(self.L, dtype=int)
        ni = 0  # not improved iterations
        while ni < self.O_max:  # inter-routing optimization
            move_gain = self.calc_move_gain(s)  # calculate the move gain
            while True:
                if np.max(move_gain) <= 0:
                    break
                k1, i1, k2, i2 = np.unravel_index(np.argmax(move_gain), move_gain.shape)  # best move
                new_s = self.perform_insertion_return_new(s, k1, i1, k2, i2)
                if H1[self.hash_func1(new_s) % self.L] == 1 and H2[self.hash_func2(new_s) % self.L] == 1:
                    move_gain[k1][i1][k2][i2] = 0  # forbidden by the tabu list, set move gain zero
                else:
                    s = self.perform_insertion(s, k1, i1, k2, i2)  # perform the best non-tabu move
                    H1[self.hash_func1(s)] = 1
                    H2[self.hash_func2(s)] = 1
                    R.add(k1)
                    R.add(k2)
                    break
            if uniform(0, 1) < self.alpha * (sum([len(s.tours[_]) for _ in R])) / self.n:
                s = self.intra_route_optimization(s)  # perform intra-routing optimization
                R.clear()
            if s.total < s_best.total:
                s_best = deepcopy(s)
                ni = 0
            else:
                ni = ni + 1
        return s_best

    def intra_route_optimization(self, s):
        for _ in range(len(s.tours)):
            if uniform(0, 1) > self.Ps:
                s.tours[_] = self.two_opt(s.tours[_])
            else:
                s.tours[_] = self.simple_tabu_search(s.tours[_])
        return self.calc_length(s)

    def calc_move_gain(self, s):
        move_gain = np.zeros(shape=(self.m, self.n + 1, self.m, self.n + 1), dtype=float)
        for k2 in range(len(s.tours)):
            for idx2 in range(1, len(s.tours[k2]) - 1):
                i2_p, i2, i2_n = s.tours[k2][idx2 - 1], s.tours[k2][idx2], s.tours[k2][idx2 + 1]
                if len(self.salesmen_of(i2)) > 1:
                    for k1 in self.salesmen_of(i2):
                        if k1 != k2:
                            for idx1 in range(0, len(s.tours[k1]) - 1):
                                i1, i1_n = s.tours[k1][idx1], s.tours[k1][idx1 + 1]
                                move_gain[k1][i1][k2][i2] = self.Distance[i1][i2] + self.Distance[i2][i1_n] \
                                                          + self.Distance[i2_p][i2_n] - self.Distance[i1][i1_n] \
                                                          - self.Distance[i2_p][i2] - self.Distance[i2][i2_n]
        return move_gain

    def perform_insertion(self, s, k1, i1, k2, i2):
        # print(f'move {i2} in tour {k2} to the position after {i1} in tour {k1}')
        s.tours[k2].remove(i2)
        idx = -1
        for i in range(len(s.tours[k1])):
            if s.tours[k1][i] == i1:
                idx = i
                break
        assert idx != -1
        s.tours[k1].insert(idx, i2)
        return self.calc_length(s)

    def perform_insertion_return_new(self, s, k1, i1, k2, i2):
        new_s = deepcopy(s)
        new_s.tours[k2].remove(i2)
        idx = -1
        for i in range(len(new_s.tours[k1])):
            if new_s.tours[k1][i] == i1:
                idx = i
                break
        assert idx != -1
        new_s.tours[k1].insert(idx, i2)
        return self.calc_length(new_s)

    def local_optima_escaping(self, s, ni):
        Pd = 1 - np.exp(-ni / self.T) / 2
        remove = []
        for r, tour in enumerate(s.tours):
            for i in range(len(tour) - 2, 0, -1):
                if uniform(0, 1) < Pd:
                    remove.append(tour[i])
        shuffle(remove)
        for city in remove:
            self.optimal_insert(s, city)
        return self.calc_length(s)

    @staticmethod
    def hash_func1(s):
        return sum([(k + 1) * sum([tour[i] for i in range(1, len(tour) - 1)]) for k, tour in enumerate(s.tours)])

    @staticmethod
    def hash_func2(s):
        return sum([sum([tour[i - 1] * tour[i] for i in range(2, len(tour) - 1)]) for tour in s.tours])

    def hash_func1_neighbor(self, s, k1, k2, i2):
        return self.hash_func1(s) + i2 * k1 - i2 * k2

    def hash_func2_neighbor(self, s, k1, i1, k2, i2):
        i1_n, i2_p, i2_n = 0, 0, 0
        for idx in range(len(s.tours[k1])):
            if s.tours[k1] == i1:
                i1_n = s.tours[k1][idx + 1]
                break
        for idx in range(len(s.tours[k2])):
            if s.tours[k2] == i2:
                i2_p, i2_n = s.tours[k2][idx + 1], s.tours[k2][idx + 1]
                break
        assert i1_n and i2_p and i2_n
        return self.hash_func2(s) + i1 * i2 + i2 * i1_n - i1 * i1_n + i2_p * i2_n - i2_p * i2 - i2 * i2_n

    def two_opt(self, tour):
        # evo = [vcmtsp.cost(tour)]
        # exchange = 0
        for _ in range(self.Opt_depth):
            _i, _j, improve = -1, -1, -1  # best 2-opt move
            for i in range(1, len(tour)):  # find the best 2-opt move
                for j in range(i + 1, len(tour)):
                    d1 = self.Distance[tour[i - 1]][tour[i]] + self.Distance[tour[j - 1]][tour[j]]  # remove edges
                    d2 = self.Distance[tour[i - 1]][tour[j - 1]] + self.Distance[tour[i]][tour[j]]  # insert edges
                    if d1 - d2 > improve:
                        _i, _j, improve = i, j, d1 - d2  # best 2-opt move
            if improve > 0:  # perform the best 2-opt move
                tour = tour[:_i] + list(reversed(tour[_i:_j])) + tour[_j:]
                # exchange += 1
                # evo.append(vcmtsp.cost(tour))
                # plt.clf()
                # plt.subplot(1, 2, 1)
                # plt.plot(range(len(evo)), evo)
                # plt.title(f'2-opt best move {exchange}: {evo[-1]:.2f}')
                # plt.subplot(1, 2, 2)
                # vcmtsp.draw(plt, routes=[tour])
                # plt.pause(0.001)
            else:
                break  # no improving edge exchange exits, break
        # plt.show()
        return tour

    def two_opt_inplace(self, tour):
        evo = [vcmtsp.cost(tour)]
        exchange = 0
        for i in range(1, len(tour)):
            for j in range(i + 1, len(tour)):
                d1 = self.Distance[tour[i - 1]][tour[i]] + self.Distance[tour[j - 1]][tour[j]]
                d2 = self.Distance[tour[i - 1]][tour[j - 1]] + self.Distance[tour[i]][tour[j]]
                if d2 < d1:
                    tour = tour[:i] + list(reversed(tour[i:j])) + tour[j:]
                    exchange += 1
                    evo.append(vcmtsp.cost(tour))
                    plt.clf()
                    plt.subplot(1, 2, 1)
                    plt.plot(range(len(evo)), evo)
                    plt.title(f'exchange inplace {exchange}: {evo[-1]:.2f}')
                    plt.subplot(1, 2, 2)
                    vcmtsp.draw(plt, routes=[tour])
                    plt.pause(0.001)
        plt.show()
        return tour

    def simple_tabu_search(self, tour):
        # evo = [vcmtsp.cost(tour)]
        # exchange = 0
        tabu = np.zeros(shape=(self.n + 1, self.n + 1), dtype=int)
        ni, tl = 0, int(self.Tl * len(tour))
        while ni < len(tour):
            _i, _j, improve = -1, -1, -1  # best 2-opt move
            for i in range(1, len(tour)):  # find the best 2-opt move
                for j in range(i + 1, len(tour)):
                    d1 = self.Distance[tour[i - 1]][tour[i]] + self.Distance[tour[j - 1]][tour[j]]  # remove edges
                    d2 = self.Distance[tour[i - 1]][tour[j - 1]] + self.Distance[tour[i]][tour[j]]  # insert edges
                    if d1 - d2 > improve and \
                            tabu[tour[i - 1]][tour[i]] == 0 and tabu[tour[j - 1]][tour[j]] == 0 and \
                            tabu[tour[i - 1]][tour[j - 1]] == 0 and tabu[tour[i]][tour[j]] == 0:
                        _i, _j, improve = i, j, d1 - d2  # best 2-opt move
            if improve > 0:
                tour = tour[:_i] + list(reversed(tour[_i:_j])) + tour[_j:]
                tabu[tour[_i - 1]][tour[_i]] = tl  # set tenure
                tabu[tour[_j - 1]][tour[_j]] = tl
                tabu[tour[_i - 1]][tour[_j - 1]] = tl
                tabu[tour[_i]][tour[_j]] = tl
                ni = 0
                # exchange += 1
                # evo.append(vcmtsp.cost(tour))
                # plt.clf()
                # plt.subplot(1, 2, 1)
                # plt.plot(range(len(evo)), evo)
                # plt.title(f'STS eligible move {exchange}: {evo[-1]:.2f}')
                # plt.subplot(1, 2, 2)
                # vcmtsp.draw(plt, routes=[tour])
                # plt.pause(0.001)
            else:
                ni = ni + 1
            tabu = tabu - 1
            tabu = np.clip(tabu, a_min=0, a_max=np.inf)
            # print(tabu)
            # print(ni)
        # plt.show()
        return tour


plt.figure(figsize=(10, 4))
case, m = 'att48', 4
vcmtsp = Instance(file=f'../../data/TSPLIB_VC/{case}.tsp/{case}.tsp.txt', m=m)
optimizer = ITPLS(data=vcmtsp)
optimizer.optimize()
# optimizer.simple_tabu_search(tour=list(range(1, 201)) + [1])
